"""
信号预处理模块

实现了信号预处理功能，包括：
- 去噪
- 归一化
- 滤波
- 重采样
"""

import numpy as np
from typing import Optional, Tuple
from loguru import logger
from scipy import signal
from .median_filter import MedianFilter


class SignalPreprocessing:
    """信号预处理器"""
    
    def __init__(self, 
                 sample_rate: int = 100,
                 target_sample_rate: Optional[int] = None,
                 normalize: bool = True,
                 denoise: bool = True,
                 filter_type: str = "median"):
        """
        初始化信号预处理器
        
        Args:
            sample_rate: 原始采样率
            target_sample_rate: 目标采样率
            normalize: 是否归一化
            denoise: 是否去噪
            filter_type: 滤波器类型
        """
        self.sample_rate = sample_rate
        self.target_sample_rate = target_sample_rate
        self.normalize = normalize
        self.denoise = denoise
        self.filter_type = filter_type
        
        # 初始化滤波器
        if filter_type == "median":
            self.filter = MedianFilter()
        else:
            self.filter = None
        
        logger.info("初始化信号预处理器")
    
    def preprocess(self, signal_data: np.ndarray) -> np.ndarray:
        """
        预处理信号
        
        Args:
            signal_data: 原始信号
            
        Returns:
            预处理后的信号
        """
        processed_signal = signal_data.copy()
        
        # 去噪
        if self.denoise and self.filter:
            processed_signal = self.filter.filter(processed_signal)
        
        # 归一化
        if self.normalize:
            processed_signal = self._normalize_signal(processed_signal)
        
        # 重采样
        if self.target_sample_rate and self.target_sample_rate != self.sample_rate:
            processed_signal = self._resample_signal(processed_signal)
        
        return processed_signal
    
    def _normalize_signal(self, signal_data: np.ndarray) -> np.ndarray:
        """
        归一化信号
        
        Args:
            signal_data: 输入信号
            
        Returns:
            归一化后的信号
        """
        # 去除直流分量
        signal_centered = signal_data - np.mean(signal_data)
        
        # 幅值归一化
        max_amplitude = np.max(np.abs(signal_centered))
        if max_amplitude > 0:
            signal_normalized = signal_centered / max_amplitude
        else:
            signal_normalized = signal_centered
        
        return signal_normalized
    
    def _resample_signal(self, signal_data: np.ndarray) -> np.ndarray:
        """
        重采样信号
        
        Args:
            signal_data: 输入信号
            
        Returns:
            重采样后的信号
        """
        if not self.target_sample_rate:
            return signal_data
        
        # 计算重采样比例
        ratio = self.target_sample_rate / self.sample_rate
        
        # 重采样
        resampled_signal = signal.resample(signal_data, int(len(signal_data) * ratio))
        
        return resampled_signal
    
    def __str__(self) -> str:
        return (f"SignalPreprocessing(sample_rate={self.sample_rate}, "
                f"normalize={self.normalize}, denoise={self.denoise})") 